{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 基础设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# 获取项目根目录的绝对路径\n",
    "current_dir = os.path.dirname(os.path.abspath(''))\n",
    "\n",
    "# 确保项目根目录在系统路径中\n",
    "if current_dir not in sys.path:\n",
    "    sys.path.append(current_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from tqdm import tqdm\n",
    "\n",
    "from student.time_series import LSTMPredictor\n",
    "from data.distill_dataset import DistillDataManager"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 加载半径矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radius_save_path = \"../data/radius_matrix.pth\"\n",
    "\n",
    "if os.path.exists(radius_save_path):\n",
    "    radius_matrix = torch.load(radius_save_path, weights_only=False)\n",
    "else:\n",
    "    radius_matrix = torch.rand(10, 10) * 10\n",
    "    torch.save(radius_matrix, radius_save_path)\n",
    "\n",
    "print(radius_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 加载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_manager = DistillDataManager(radius_matrix)\n",
    "\n",
    "# 加载训练数据集并分割出验证集\n",
    "full_train_dataset = data_manager.get_train_dataset(mode='serial')\n",
    "train_size = int(0.8 * len(full_train_dataset))\n",
    "val_size = len(full_train_dataset) - train_size\n",
    "train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])\n",
    "\n",
    "# 加载测试数据集\n",
    "test_dataset = data_manager.get_test_dataset(mode='serial')\n",
    "\n",
    "print(\"训练集大小:\", len(train_dataset))\n",
    "print(\"验证集大小:\", len(val_dataset))\n",
    "print(\"测试集大小:\", len(test_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型训练\n",
    "\n",
    "## 定义和初始化模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"当前GPU内存使用: {torch.cuda.memory_allocated(0)/1024**2:.1f} MB\")\n",
    "\n",
    "model = LSTMPredictor(\n",
    "    input_dim=11,  # 10个端口 + 1个零通道\n",
    "    hidden_dim=256,\n",
    "    output_dim=10,\n",
    "    num_layers=2,\n",
    "    dropout=0.1\n",
    ").to(device)\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"模型总参数量: {total_params:,}\")\n",
    "\n",
    "model = data_manager.load_model(model, mode='serial')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 100000\n",
    "learning_rate = 1e-3\n",
    "batch_size = 32\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)\n",
    "criterion = torch.nn.MSELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_losses = []\n",
    "val_losses = []\n",
    "best_val_loss = float('inf')\n",
    "best_state_dict = None\n",
    "best_epoch = 0\n",
    "\n",
    "progress_bar = tqdm(total=num_epochs)\n",
    "\n",
    "try:\n",
    "    for epoch in range(num_epochs):\n",
    "        # 训练阶段\n",
    "        model.train()\n",
    "        epoch_loss = 0\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "        epoch_loss /= len(train_loader)\n",
    "        train_losses.append(epoch_loss)\n",
    "\n",
    "        # 验证阶段\n",
    "        model.eval()\n",
    "        val_loss = 0\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in val_loader:\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "                val_loss += criterion(outputs, labels).item()\n",
    "        \n",
    "        val_loss /= len(val_loader)\n",
    "        val_losses.append(val_loss)\n",
    "        \n",
    "        # 学习率调整\n",
    "        scheduler.step(val_loss)\n",
    "\n",
    "        # 保存最佳模型\n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            best_state_dict = model.state_dict()\n",
    "            best_epoch = epoch\n",
    "\n",
    "        # 更新进度条\n",
    "        current_lr = optimizer.param_groups[0]['lr']\n",
    "        progress_bar.set_description(\n",
    "            f\"Epoch {epoch+1}, Train Loss: {epoch_loss:.4f}, Val Loss: {val_loss:.4f}, LR: {current_lr:.2e}\"\n",
    "        )\n",
    "        progress_bar.update(1)\n",
    "\n",
    "        # 每100个epoch保存checkpoint\n",
    "        if epoch % 100 == 0:\n",
    "            checkpoint_path = f'checkpoint_epoch_{epoch}.pth'\n",
    "            torch.save({\n",
    "                'epoch': epoch,\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'train_loss': epoch_loss,\n",
    "                'val_loss': val_loss,\n",
    "                'best_val_loss': best_val_loss,\n",
    "                'best_epoch': best_epoch\n",
    "            }, checkpoint_path)\n",
    "            print(f\"\\nCheckpoint saved: {checkpoint_path}\")\n",
    "\n",
    "        # 打印GPU内存使用\n",
    "        if torch.cuda.is_available() and epoch % 10 == 0:\n",
    "            print(f\"\\nGPU内存使用: {torch.cuda.memory_allocated(0)/1024**2:.1f} MB\")\n",
    "\n",
    "except KeyboardInterrupt:\n",
    "    print(\"\\n训练被手动中断\")\n",
    "\n",
    "finally:\n",
    "    progress_bar.close()\n",
    "    print(f\"最佳Epoch: {best_epoch}; 最佳验证损失: {best_val_loss:.4f}\")\n",
    "    \n",
    "    if best_state_dict is not None:\n",
    "        model.load_state_dict(best_state_dict)\n",
    "        save_path = data_manager.save_model(model, mode='serial')\n",
    "        print(f\"最佳模型已保存至: {save_path}\")\n",
    "\n",
    "    # 绘制损失曲线\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.plot(train_losses, label='训练损失')\n",
    "    plt.plot(val_losses, label='验证损失')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('训练和验证损失曲线')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型测试\n",
    "\n",
    "## 加载模型\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = LSTMPredictor(\n",
    "    input_dim=11,\n",
    "    hidden_dim=256,\n",
    "    output_dim=10,\n",
    "    num_layers=2,\n",
    "    dropout=0.1\n",
    ").to(device)\n",
    "\n",
    "model = data_manager.load_model(model, mode='serial')\n",
    "print(\"模型加载完成\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "batch_size = 32\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)  # 不打乱测试集\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "test_loss = 0\n",
    "all_preds = []\n",
    "all_labels = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for inputs, labels in tqdm(test_loader, desc=\"测试进度\"):\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        test_loss += loss.item()\n",
    "        all_preds.append(outputs.cpu().numpy())\n",
    "        all_labels.append(labels.cpu().numpy())\n",
    "\n",
    "test_loss /= len(test_loader)\n",
    "print(f\"测试损失: {test_loss:.6f}\")\n",
    "\n",
    "all_preds = np.concatenate(all_preds, axis=0)\n",
    "all_labels = np.concatenate(all_labels, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 可视化结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "\n",
    "# 1. 时序预测对比图\n",
    "plt.subplot(131)\n",
    "sample_idx = 0  # 可以改变这个索引来查看不同的样本\n",
    "port_idx = 0    # 可以改变这个索引来查看不同的端口\n",
    "plt.plot(all_preds[sample_idx, :, port_idx], 'b-', label='预测值', alpha=0.7)\n",
    "plt.plot(all_labels[sample_idx, :, port_idx], 'r-', label='真实值', alpha=0.7)\n",
    "plt.title(f'样本 {sample_idx}, 端口 {port_idx} 的时序预测')\n",
    "plt.xlabel('时间步')\n",
    "plt.ylabel('强度')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "\n",
    "# 2. 预测值与真实值的散点图\n",
    "plt.subplot(132)\n",
    "plt.scatter(all_labels.flatten(), all_preds.flatten(), alpha=0.01)  # 降低alpha值\n",
    "plt.plot([all_labels.min(), all_labels.max()], \n",
    "        [all_labels.min(), all_labels.max()], 'r--')  # 理想的对角线\n",
    "plt.title('预测值 vs 真实值')\n",
    "plt.xlabel('真实值')\n",
    "plt.ylabel('预测值')\n",
    "plt.grid(True)\n",
    "\n",
    "# 3. 每个端口的误差箱型图\n",
    "plt.subplot(133)\n",
    "errors = all_preds - all_labels\n",
    "plt.boxplot([errors[:,:,i].flatten() for i in range(errors.shape[2])])\n",
    "plt.title('各端口预测误差分布')\n",
    "plt.xlabel('端口')\n",
    "plt.ylabel('误差')\n",
    "plt.grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 计算性能指标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 整体性能指标\n",
    "mse = np.mean((all_preds - all_labels) ** 2)\n",
    "mae = np.mean(np.abs(all_preds - all_labels))\n",
    "rmse = np.sqrt(mse)\n",
    "r2 = 1 - np.sum((all_preds - all_labels) ** 2) / np.sum((all_labels - np.mean(all_labels)) ** 2)\n",
    "\n",
    "print(\"整体性能指标:\")\n",
    "print(f\"均方误差 (MSE): {mse:.6f}\")\n",
    "print(f\"均方根误差 (RMSE): {rmse:.6f}\")\n",
    "print(f\"平均绝对误差 (MAE): {mae:.6f}\")\n",
    "print(f\"R² 分数: {r2:.6f}\")\n",
    "\n",
    "# 时序维度的性能指标\n",
    "temporal_mse = np.mean((all_preds - all_labels) ** 2, axis=(0, 2))\n",
    "plt.figure(figsize=(10, 4))\n",
    "plt.plot(temporal_mse)\n",
    "plt.title('不同时间步的预测误差')\n",
    "plt.xlabel('时间步')\n",
    "plt.ylabel('MSE')\n",
    "plt.grid(True)\n",
    "plt.show()\n",
    "\n",
    "# 各端口性能指标\n",
    "print(\"\\n各端口性能指标:\")\n",
    "for port in range(all_preds.shape[2]):\n",
    "    port_preds = all_preds[:,:,port]\n",
    "    port_labels = all_labels[:,:,port]\n",
    "    port_mse = np.mean((port_preds - port_labels) ** 2)\n",
    "    port_mae = np.mean(np.abs(port_preds - port_labels))\n",
    "    port_rmse = np.sqrt(port_mse)\n",
    "    port_r2 = 1 - np.sum((port_preds - port_labels) ** 2) / \\\n",
    "              np.sum((port_labels - np.mean(port_labels)) ** 2)\n",
    "    print(f\"\\n端口 {port}:\")\n",
    "    print(f\"  MSE: {port_mse:.6f}\")\n",
    "    print(f\"  RMSE: {port_rmse:.6f}\")\n",
    "    print(f\"  MAE: {port_mae:.6f}\")\n",
    "    print(f\"  R²: {port_r2:.6f}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
